package com.aboat.model;


import cn.edu.hfut.dmic.webcollector.model.CrawlDatums;
import cn.edu.hfut.dmic.webcollector.model.Links;
import cn.edu.hfut.dmic.webcollector.model.Page;

import cn.edu.hfut.dmic.webcollector.plugin.berkeley.BreadthCrawler;
import cn.edu.hfut.dmic.webcollector.util.FileUtils;
import com.aboat.view.WindowFrame;
import org.jsoup.nodes.Element;
import org.jsoup.select.Elements;


import java.io.File;
import java.io.IOException;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 抓取图片
 *
 * @author aboat
 * @email clear_windy@yeah.net
 * @create 2017-03-04 22:39
 */
public class ImageCrawler extends BreadthCrawler {
    //AtomicInteger imageId = new AtomicInteger(0);

    //原子性int，用于生成图片文件名
    AtomicInteger imageId;

    //用于保存图片的文件夹
    File downloadDir;

    String imgTag;//图片标签

    public ImageCrawler(String crawlPath,String downloadPath) {
        super(crawlPath, true);
        downloadDir = new File(downloadPath);
        if(!downloadDir.exists()){
            downloadDir.mkdirs();
        }
        computeImageId();
    }

    @Override
    public void visit(Page page, CrawlDatums next) {
        //根据http头中的Content-Type信息来判断当前资源是网页还是图片
        String contentType = page.getResponse().getContentType();
        if(contentType==null){
            return;
        }else if (contentType.contains("html")) {
            //如果是网页，则抽取其中包含图片的URL，放入后续任务
            Elements imgs = page.select(imgTag);
            //Elements imgs = page.getDoc().select("input[src]");
            for (Element img : imgs) {
                String imgSrc = img.attr("abs:src");
                if(imgSrc.indexOf("thumb")<0) {
                    next.add(imgSrc);
                }
            }
        } else if (contentType.startsWith("image")) {
            if(WindowFrame.maxNum != null) {
                if (imageId.get() >= WindowFrame.maxNum) {
                    WindowFrame.textContent.append("未注册用户只能爬取" + WindowFrame.maxNum + "张图片" + "\r\n");
                    WindowFrame.textContent.selectAll();
                    this.stop();
                    this.setResumable(false);
                }
            }else if(imageId.get()%500 == 0){
                WindowFrame.textContent.setText("");
            }

            //如果是图片，直接下载
            String extensionName=contentType.split("/")[1];
            String imageFileName=imageId.incrementAndGet()+"."+extensionName;
            File imageFile=new File(downloadDir,imageFileName);
            try {
                FileUtils.writeFile(imageFile, page.getContent());
                WindowFrame.textContent.append("保存图片:"+page.getUrl()+"\r\n");
                WindowFrame.textContent.selectAll();
            } catch (IOException ex) {
                throw new RuntimeException(ex);
            }
        }
    }

    public void computeImageId(){
        int maxId=-1;
        for(File imageFile:downloadDir.listFiles()){
            String fileName=imageFile.getName();
            String idStr=fileName.split("\\.")[0];
            int id=Integer.valueOf(idStr);
            if(id>maxId){
                maxId=id;
            }
        }
        imageId=new AtomicInteger(maxId);
    }

    public String getImgTag() {
        return imgTag;
    }

    public void setImgTag(String imgTag) {
        this.imgTag = imgTag;
    }

    public AtomicInteger getImageId() {
        return imageId;
    }

    public void setImageId(AtomicInteger imageId) {
        this.imageId = imageId;
    }
}
